/*
 *  Extras.cpp
 *  vorpal
 *
 *  Created by Russell Y Webb on 1/10/09.
 *  Copyright 2009 UC. All rights reserved.
 *
 */

#include "Extras.h"

#include "Common.h"
#include "Messages.h"
#include "VmGroup.h"
#include <math.h>

using namespace Vorpal;

class Installer{
public:
	Installer(ReceiveHook hook){ 
		TheGroup().AddReceiveHook(hook);
	}
};

// ------------------------------------------------------------------------------------------------
// --------------------- Pseudo Random Number Generator -------------------------------------------
// ------------------------------------------------------------------------------------------------

/* Period parameters */  
#define N 624
#define M 397
#define MATRIX_A 0x9908b0df   /* constant vector a */
#define UPPER_MASK 0x80000000 /* most significant w-r bits */
#define LOWER_MASK 0x7fffffff /* least significant r bits */

/* Tempering parameters */   
#define TEMPERING_MASK_B 0x9d2c5680
#define TEMPERING_MASK_C 0xefc60000
#define TEMPERING_SHIFT_U(y)  (y >> 11)
#define TEMPERING_SHIFT_S(y)  (y << 7)
#define TEMPERING_SHIFT_T(y)  (y << 15)
#define TEMPERING_SHIFT_L(y)  (y >> 18)

/** 
 \brief Mersenne Twister random number generator
 
 This library is free software under the Artistic license.       
 Copyright (C) 1997, 1999 Makoto Matsumoto and Takuji Nishimura. 
 Any feedback is very welcome. For any question, comments,       
 see http://www.math.keio.ac.jp/matumoto/emt.html or email       
 matumoto@math.keio.ac.jp                                        
 */
class RandomSource{
public:
	/// \param seed if zero then seconds since 1970 is used
    RandomSource(uint32_t seed = 0) : initialSeed(seed), mti(N+1){ 
		mag01[0] = 0x0;
		mag01[1] = MATRIX_A;
		
		if(initialSeed == 0){
			initialSeed = time(0);
			//initialSeed = clock();
			if(initialSeed == 0){
				initialSeed = 0xFA43AB56;
			}
		}
		sgenrand(initialSeed);
	}
    ~RandomSource(){ }
	
	
	/// \return a random 32-bit unsigned integer
	uint32_t UInt32(){
		uint32_t y;
		/* mag01[x] = x * MATRIX_A  for x=0,1 */
		
		if (mti >= N) { /* generate N words at one time */
			int kk;
			
			if (mti == N+1)   /* if sgenrand() has not been called, */
				sgenrand(4357); /* a default initial seed is used   */
			
			for (kk=0;kk<N-M;kk++) {
				y = (mt[kk]&UPPER_MASK)|(mt[kk+1]&LOWER_MASK);
				mt[kk] = mt[kk+M] ^ (y >> 1) ^ mag01[y & 0x1];
			}
			for (;kk<N-1;kk++) {
				y = (mt[kk]&UPPER_MASK)|(mt[kk+1]&LOWER_MASK);
				mt[kk] = mt[kk+(M-N)] ^ (y >> 1) ^ mag01[y & 0x1];
			}
			y = (mt[N-1]&UPPER_MASK)|(mt[0]&LOWER_MASK);
			mt[N-1] = mt[M-1] ^ (y >> 1) ^ mag01[y & 0x1];
			
			mti = 0;
		}
		
		y = mt[mti++];
		y ^= TEMPERING_SHIFT_U(y);
		y ^= TEMPERING_SHIFT_S(y) & TEMPERING_MASK_B;
		y ^= TEMPERING_SHIFT_T(y) & TEMPERING_MASK_C;
		y ^= TEMPERING_SHIFT_L(y);
		
		return y;     
	}
    
	/// \return a value from a Normal distribution (mean = 0, var = 1)
	double Normal(){
		double u = Float();
		double v = Float();
		// half of Box-Muller algorithm (could get another sample from u and v)
		return sqrt(-2 * log(u)) * cos(2 * M_PI * v);
	}
	
	/// \return a random doublein range [0, 1)
	double Float(){
		// [0, 1]
		return ((double)UInt32()) * 2.3283064370807974e-10;
	}
	
	/// \return the seed that can be used to recreate this RandomSource
	uint32_t Seed() const{ return initialSeed; }
private:
    void sgenrand(uint32_t seed){
		for(int i=0;i<N;i++) {
			mt[i] = seed & 0xffff0000;
			seed = 69069 * seed + 1;
			mt[i] |= (seed & 0xffff0000) >> 16;
			seed = 69069 * seed + 1;
		}
		mti = N;	
	}
	
	uint32_t initialSeed;
    uint32_t mt[N]; /* the array for the state vector  */
    int mti; /* mti==N+1 means mt[N] is not initialized */
    uint32_t mag01[2];
};

RandomSource &R(){
	static RandomSource r;
	return r;
}

class RandomGenerator : public Unique, public RandomSource{
public:
	RandomGenerator(uint seed) : Unique("RandomGenerator", false), RandomSource(seed){ }
	
	bool Receive(Message message, uint32_t numArgs, VM *vm){
		if(message == msg("uniform")){
			double result = Float();
			vm->Drop(numArgs);
			vm->PushPtr( new Number(result) );
		}
		else if(message == msg("bool")){
			bool result = UInt32() & 1;
			vm->Drop(numArgs);
			vm->Push( vm->FromBool(result) );			
		}
		else if(message == msg("normal")){
			double result = Normal();
			vm->Drop(numArgs);
			vm->PushPtr( new Number(result) );			
		}
		else{
			return Object::Receive(message, numArgs, vm);
		}
		return true;
	}
	
private:
	static Installer installer;
	
	static bool hook(Message message, uint32_t numArgs, VM *vm){
		if(message == msg("mt_random")){
			if(numArgs > 2){ Failure(ERR_INVALID_NUM_ARGS); }
			double seed = 0;
			if(numArgs == 1){
				seed = vm->FromTop(0)->ToNumber();
			}
			vm->Drop(numArgs);
			vm->PushPtr( new RandomGenerator(seed) );
			return true;
		}
		return false;
	}	
};
Installer RandomGenerator::installer(hook);


// ---------------- ESN ------------------------------
template< class Type >
class Grid{ // 2D array of doubles
public:
	Grid(uint firstDim, uint secondDim = 1) : firstDim(firstDim), secondDim(secondDim){
		Assert(firstDim > 0);
		Assert(secondDim > 0);
		size = firstDim * secondDim;
		array = new Type[size];
		Assert(array);
		loop(i, 0, size){ array[i] = 0; }
	}
	
	~Grid(){
		delete[] array;
	}
	
	Type &operator()(uint index1, uint index2 = 0){
		uint index = index1*secondDim + index2;
		Assert(index < size);
		return array[index];
	}
private:
	uint size;
	uint firstDim;
	uint secondDim;
	Type *array;
};

class Esn{
public:
	Esn(uint seed, uint numInputs, uint numLinks, uint numNeurons = 40) :
		numInputs(numInputs),
		numLinks(numLinks),
		numNeurons(numNeurons),
		firstIsCurrent(true),
		esnState1(numNeurons + numInputs),
		esnState2(numNeurons + numInputs),
		inWeights(numInputs, numNeurons),
		outWeights(numNeurons + numInputs),
		linkIndex(numNeurons, numLinks),
		innerW(numNeurons, numLinks)
	{	
		RandomSource r(seed);
		loop(ni, 0, numNeurons + numInputs){
			outWeights(ni) = (2*r.Float() - 1);
		}
		double inputLimit = r.Float() / numNeurons;
		loop(ii, 0, numInputs){
			loop(ni, 0, numNeurons){
				inWeights(ii, ni) = inputLimit * (2*r.Float() - 1);
			}
		}
		double innerLimit = 0.05 * r.Float();
		loop(ci, 0, numLinks){
			loop(ni, 0, numNeurons){
				linkIndex(ni, ci) = r.UInt32() % numNeurons;
				innerW(ni, ci) = innerLimit * (2*r.Float() - 1);
			}
		}
	}


	double Step(const double *input, const double *target = 0){ 
		Grid<double> &current = firstIsCurrent ? esnState1 : esnState2 ;
		Grid<double> &next = firstIsCurrent ? esnState2 : esnState1 ;
		firstIsCurrent = !firstIsCurrent;
		// inputs
		loop(ii, 0, numInputs){
			loop(n, 0, numNeurons){
				next(n) = current(n) + inWeights(ii, n)*input[ii];
			}
		}
		// internal
		loop(n, 0, numNeurons){
			loop(c, 0, numLinks){
				next(linkIndex(n, c)) += innerW(n, c) * current(n);
			}
		}
		loop(n, 0, numNeurons){
			next(n) = tanh(next(n));
		}
		// store inputs
		loop(i, 0, numInputs){
			next(numNeurons + i) = input[i];
		}		
		// output
		double result = 0;
		loop(n, 0, numNeurons + numInputs){
			result += outWeights(n)*next(n);
		}
		if(target){
			double trainingErr = result - atanh(*target);
			double learn = (0.10 / numNeurons) + (0.20 / numNeurons) * R().Float();
			loop(n, 0, numNeurons + numInputs){
				outWeights(n) -= learn * trainingErr * next(n);
			}
		}
		return tanh(result);
	}

protected:
	uint numInputs;
	uint numLinks;
	uint numNeurons;
	bool firstIsCurrent;
	Grid<double> esnState1;
	Grid<double> esnState2;
	Grid<double> inWeights;
	Grid<double> outWeights;
	Grid<int> linkIndex;
	Grid<double> innerW;
};

void ExtractDoubleArgs(double *args, uint firstArgIndex, uint numToExtract, uint numArgsOnStack, VM *vm){
	if(firstArgIndex + numToExtract > numArgsOnStack){ Failure(ERR_INVALID_NUM_ARGS); }
	uint argStart = vm->Depth() - numArgsOnStack;
	loop(i, 0, numToExtract){
		args[i] = Translate<Number>(vm->FromBase(argStart + firstArgIndex + i).Ptr())->Value();
	}	
}
				 
class EsnObject : public Unique, public Esn{
public:
	EsnObject(uint seed, uint numInputs, uint numLinks, uint numNeurons) : Unique("Esn", false), Esn(seed, numInputs, numLinks, numNeurons){ }
	
	bool Receive(Message message, uint32_t numArgs, VM *vm){
		if(message == msg("step")){
			if(numArgs != 2 + numInputs && numArgs != 1 + numInputs){ Failure(ERR_INVALID_NUM_ARGS); }
			uint numVals = numArgs - 1;
			DeclareArray(vals, double,  numVals);
			ExtractDoubleArgs(vals, 1, numVals, numArgs, vm);
			double *target = numVals > numInputs ? &vals[numVals - 1] : 0;
			vm->DropThenReplaceTop(numArgs - 1, new Number(Step(vals, target)) );
		}
		else{
			return Object::Receive(message, numArgs, vm);
		}
		return true;
	}
	
private:
	static Installer installer;
	
	static bool hook(Message message, uint32_t numArgs, VM *vm){
		if(message == msg("esn")){
			double vals[4];
			ExtractDoubleArgs(vals, 1, 4, numArgs, vm);
			vm->DropThenReplaceTop(numArgs - 1, new EsnObject(vals[0], vals[1], vals[2], vals[3]));
			return true;
		}
		return false;
	}	
};
Installer EsnObject::installer(hook);

